import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class WordSenseCount {
  private HashMap<String, HashMap<String, Integer>> wordSenseMap_;
  private HashMap<String, String> mostFreqSense_;
  private HashMap<String, Integer> wordCount_;
  
  public WordSenseCount() {
    wordSenseMap_ = new HashMap<String, HashMap<String,Integer>>();
    mostFreqSense_ = new HashMap<String, String>();
    wordCount_ = new HashMap<String, Integer>();
  }

  private HashMap<String, Integer> getSenseCountForWord(String word) {
    HashMap<String, Integer> hm = wordSenseMap_.get(word);
    
    if (hm == null) {
      hm = new HashMap<String, Integer>();
      wordSenseMap_.put(word, hm);
    }
    
    return hm;
  }
  
  public void addSenseForWord(String word, String sense) {
    HashMap<String, Integer> hm = getSenseCountForWord(word);
    
    Integer senseCount = hm.get(sense);
    
    if (senseCount == null) {
      hm.put(sense, 1);
    } else {
      hm.put(sense, senseCount + 1);
    }
  }
  
  public int getSenseCountForWord(String word, String sense) {
    HashMap<String, Integer> hm = getSenseCountForWord(word);
    
    Integer senseCount = hm.get(sense);
    
    if (senseCount == null) return 0;
    
    return senseCount;
  }
  
  public void processLexelt(String word) {
    HashMap<String, Integer> hm = wordSenseMap_.get(word);
    if (hm == null) {
      System.err.println("Processing max sense for not existing word in training.");
      return;
    }

    Iterator<Map.Entry<String, Integer>> it = wordSenseMap_.get(word).entrySet().iterator();
    Integer maxCount = -1;
    String bestSense = null;
    Integer totSenseCount = 0; 
    while (it.hasNext()) {
      Map.Entry<String, Integer> mp = it.next();
      totSenseCount += mp.getValue();
      
      if (mp.getValue() > maxCount) {
        bestSense = mp.getKey();        
      }
    }
    
    mostFreqSense_.put(word, bestSense);
    wordCount_.put(word, totSenseCount);
  }
  
  public String getMostFrequentSense(String word) {
    return mostFreqSense_.get(word);
  }
  
  public int getWordCount(String word) {
    Integer count = wordCount_.get(word);
    if (count == null) return 0;
    return (int)count;
  }
}
